import pyrtl
import or1200_definitions

width = or1200_definitions.OR1200_OPERAND_WIDTH

class Or1200Reg2Mem(object):
    def __init__(self):

        # I/O
        self.addr = pyrtl.WireVector(bitwidth=2)
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.regdata = pyrtl.WireVector(bitwidth=width)
        # self.addr = pyrtl.Input(bitwidth=2)
        # self.lsu_op = pyrtl.Input(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        # self.regdata = pyrtl.Input(bitwidth=width)
        self.memdata = pyrtl.WireVector(bitwidth=width)

        # Internal regs and wires
        self.memdata_hh = pyrtl.WireVector(bitwidth=8, name='memdata_hh')
        self.memdata_hl = pyrtl.WireVector(bitwidth=8, name='memdata_hl')
        self.memdata_lh = pyrtl.WireVector(bitwidth=8, name='memdata_lh')
        self.memdata_ll = pyrtl.WireVector(bitwidth=8, name='memdata_ll')
        self.memdata <<= pyrtl.concat(self.memdata_hh, self.memdata_hl,self.memdata_lh, \
                                    self.memdata_ll)

        #initialize all the parts
        mux2Memdata_1 = Mux2Memdata_1()
        mux2Memdata_2 = Mux2Memdata_2()
        mux2Memdata_3 = Mux2Memdata_3()
        # mux2Memdata4 = Mux2Memdata_4()

        #establish connection relations for Mux to memdata[24: ]
        self.memdata_hh <<= mux2Memdata_1.memdata_hh
        mux2Memdata_1.lsu_op <<= self.lsu_op
        mux2Memdata_1.addr <<= self.addr
        mux2Memdata_1.regdata <<= self.regdata

        # establish connection relations for Mux to memdata[16:24]
        self.memdata_hl <<= mux2Memdata_2.memdata_hl
        mux2Memdata_2.lsu_op <<= self.lsu_op
        mux2Memdata_2.addr <<= self.addr
        mux2Memdata_2.regdata <<= self.regdata

        # establish connection relations for Mux to memdata[8:16]
        self.memdata_lh <<= mux2Memdata_3.memdata_lh
        mux2Memdata_3.lsu_op <<= self.lsu_op
        mux2Memdata_3.addr <<= self.addr
        mux2Memdata_3.regdata <<= self.regdata

        # establish connection relations for Mux to memdata[7:0]
        # self.memdata_ll <<= mux2Memdata4.memdata_ll
        # mux2Memdata4.regdata <<= self.regdata
        self.memdata_ll <<= self.regdata[0:8]


# Mux to memdata[24: ]
class Mux2Memdata_1(object):
    def __init__(self):
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.addr = pyrtl.WireVector(bitwidth=2)
        self.regdata = pyrtl.WireVector(bitwidth=width)

        self.memdata_hh = pyrtl.Register(bitwidth=8)

        with pyrtl.conditional_assignment:
            with pyrtl.concat(self.lsu_op, self.addr[0:2]) == \
                 pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b0, bitwidth=2)):
                self.memdata_hh.next |= self.regdata[0:8]
            with pyrtl.concat(self.lsu_op, self.addr[0:2]) == \
                 pyrtl.concat(or1200_definitions.OR1200_LSUOP_SH, pyrtl.Const(0b0, bitwidth=2)):
                self.memdata_hh.next |= self.regdata[8:16]
            with pyrtl.otherwise:
                self.memdata_hh.next |= self.regdata[24:32]


# Mux to memdata[16:24]
class Mux2Memdata_2(object):
    def __init__(self):
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.addr = pyrtl.WireVector(bitwidth=2)
        self.regdata = pyrtl.WireVector(bitwidth=width)

        self.memdata_hl = pyrtl.Register(bitwidth=8)

        with pyrtl.conditional_assignment:
            with pyrtl.concat(self.lsu_op, self.addr[0:2]) == \
                 pyrtl.concat(or1200_definitions.OR1200_LSUOP_SW, pyrtl.Const(0b0, bitwidth=2)):
                self.memdata_hl.next |= self.regdata[16:24]
            with pyrtl.otherwise:
                self.memdata_hl.next |= self.regdata[0:8]


# Mux to memdata[8:16]
class Mux2Memdata_3(object):
    def __init__(self):
        self.lsu_op = pyrtl.WireVector(bitwidth=or1200_definitions.OR1200_LSUOP_WIDTH)
        self.addr = pyrtl.WireVector(bitwidth=2)
        self.regdata = pyrtl.WireVector(bitwidth=width)

        self.memdata_lh = pyrtl.Register(bitwidth=8)

        with pyrtl.conditional_assignment:
            with pyrtl.concat(self.lsu_op, self.addr[0:2]) == \
                 pyrtl.concat(or1200_definitions.OR1200_LSUOP_SB, pyrtl.Const(0b0, bitwidth=2)):
                self.memdata_lh.next |= self.regdata[0:8]
            with pyrtl.otherwise:
                self.memdata_lh.next |= self.regdata[8:16]

# Mux to memdata[7:0]
# class Mux2Memdata_4(object):
#     def __init__(self):
#         self.regdata = pyrtl.WireVector(bitwidth=width, name='regdata')
#         self.memdata_ll = pyrtl.Register(bitwidth=8, name='memdata_ll')
#
#         with pyrtl.conditional_assignment:
#             self.memdata_ll.next |= self.regdata[0:8]


if __name__ == '__main__':
    or1200Reg2Men = Or1200Reg2Mem()
    print(pyrtl.working_block())
    ff = open("D:\work_doc\PyRTL\examples\Opencore\or1200_reg2mem.c", 'w')
    essent = pyrtl.ESSENT()
    essent._create_code(lambda s: ff.write(s + '\n'))
    ff.close()
    com = pyrtl.CompiledSimulation()

    input_width = []
    from pyrtl.wire import Input

    IN = pyrtl.working_block().wirevector_subset(Input)
    for inn in IN:
        name = inn.name
        width = inn.bitwidth
        input_width.append([name, width])

    import random

    for i in range(10):

        inputs = {}
        inputs_old = {}

        for inp in input_width:
            inputs[inp[0]] = []

            result = random.randint(0, 2 ** inp[1] - 1)
            result_old = result
            inputs_old[inp[0]] = [result, result_old]

        for j in range(10000):
            for inp in input_width:
                if random.randint(0, 20) in [0, 5]:
                    result = random.randint(0, 2 ** inp[1] - 1)
                    inputs[inp[0]].append(result)
                    inputs_old[inp[0]][1] = result
                else:
                    inputs[inp[0]].append(inputs_old[inp[0]][1])

        time1 = essent.step_multiple(inputs)
        time2 = com.step_multiple(inputs)

        print('==========')
        print('essent:', time1)
        print('com:', time2)
        print(time2 / time1)